#include <cusp/detail/device/spmm/coo.h>

namespace cusp {
namespace detail {
namespace device {

template <typename Matrix1,
          typename Matrix2,
          typename Matrix3,
          typename Matrix4>
void spmmm_coo(const Matrix1 &A,
               const Matrix2 &B,
               const Matrix3 &C,
                     Matrix4 &D)
{
  typedef typename Matrix4::index_type   IndexType;
  typedef typename Matrix4::value_type   ValueType;
  typedef typename Matrix4::memory_space MemorySpace;

  // check whether matrices are empty
  if (A.num_entries == 0 || B.num_entries == 0 || C.num_entries == 0)
  {
    D.resize(A.num_rows, C.num_cols, 0);
    return;
  }

  // compute row offsets for B
  cusp::array1d<IndexType,MemorySpace> B_row_offsets(B.num_rows + 1);
  cusp::detail::indices_to_offsets(B.row_indices, B_row_offsets);

  // compute row lengths for B
  cusp::array1d<IndexType,MemorySpace> B_row_lengths(B.num_rows);
  thrust::transform(B_row_offsets.begin() + 1, B_row_offsets.end(), B_row_offsets.begin(), B_row_lengths.begin(), \
      thrust::minus<IndexType>());

  // for each element A(i,j) compute the number of nonzero elements in B(j,:)
  cusp::array1d<IndexType,MemorySpace> segment_lengths(A.num_entries);
  thrust::gather(A.column_indices.begin(), A.column_indices.end(),
      B_row_lengths.begin(),
      segment_lengths.begin());

  // output pointer
  cusp::array1d<IndexType,MemorySpace> output_ptr(A.num_entries + 1);
  thrust::exclusive_scan(segment_lengths.begin(), segment_lengths.end(),
      output_ptr.begin(),
      IndexType(0));
  output_ptr[A.num_entries] = output_ptr[A.num_entries - 1] + segment_lengths[A.num_entries - 1]; // XXX is this necessary?

  // number of non-zeros in result?
  size_t coo_num_nonzeros = output_ptr[A.num_entries];

  size_t workspace_capacity = thrust::min<size_t>(coo_num_nonzeros, 16 << 20);

  {
    // TODO abstract this
    size_t free, total;
    cudaMemGetInfo(&free, &total);

    // divide free bytes by the size of each workspace unit
    size_t max_workspace_capacity = free / (4 * sizeof(IndexType) + sizeof(ValueType));

    // use at most one third of the remaining capacity
    workspace_capacity = thrust::min<size_t>(max_workspace_capacity / 3, workspace_capacity);
  }

  // workspace arrays
  cusp::array1d<IndexType,MemorySpace> A_gather_locations;
  cusp::array1d<IndexType,MemorySpace> B_gather_locations;
  cusp::array1d<IndexType,MemorySpace> I;
  cusp::array1d<IndexType,MemorySpace> J;
  cusp::array1d<ValueType,MemorySpace> V;

  size_t C_begin_row;
  size_t C_end_row;
  size_t C_begin_segment;
  size_t C_end_segment;

  cusp::array1d<IndexType,MemorySpace> temp_gather_locations;
  cusp::array1d<IndexType,MemorySpace> C_gather_locations;
  cusp::array1d<IndexType,MemorySpace> C_I;
  cusp::array1d<IndexType,MemorySpace> C_J;
  cusp::array1d<ValueType,MemorySpace> C_V;

  // decompose C = A * B into several C[slice,:] = A[slice,:] * B operations
  typedef typename cusp::coo_matrix<IndexType,ValueType,MemorySpace> Container;
  typedef typename std::list<Container> ContainerList;

  // storage for C[slice,:] partial results
  ContainerList slices;

  // compute row offsets for A
  cusp::array1d<IndexType,MemorySpace> A_row_offsets(A.num_rows + 1);
  cusp::detail::indices_to_offsets(A.row_indices, A_row_offsets);

  // compute workspace requirements for each row
  cusp::array1d<IndexType,MemorySpace> cummulative_row_workspace(A.num_rows);
  thrust::gather(A_row_offsets.begin() + 1, A_row_offsets.end(),
      output_ptr.begin(),
      cummulative_row_workspace.begin());

  size_t begin_row = 0;
  size_t total_work = 0;
  size_t call_no = 0;

  while (begin_row < size_t(A.num_rows))
  {
    Container temp_slice;
    Container R_slice;

    // find largest end_row such that the capacity of [begin_row, end_row) fits in the workspace_capacity
    size_t end_row = thrust::upper_bound(cummulative_row_workspace.begin() + begin_row, cummulative_row_workspace.end(),
        total_work + IndexType(workspace_capacity)) - cummulative_row_workspace.begin();
    // size_t end_row = begin_row + 20;

    size_t begin_segment = A_row_offsets[begin_row];
    size_t end_segment   = A_row_offsets[end_row];

    // TODO throw exception signaling that there is insufficient memory (not necessarily bad_alloc)
    //if (begin_row == end_row)
    //    // workspace wasn't large enough, throw cusp::memory_allocation_failure?

    size_t workspace_size = output_ptr[end_segment] - output_ptr[begin_segment];

    total_work += workspace_size;

    // TODO remove these when an exception is in place
    assert(end_row > begin_row);
    assert(workspace_size <= workspace_capacity);

    coo_spmm_helper(workspace_size,
        begin_row, end_row,
        begin_segment, end_segment,
        A, B, temp_slice,
        B_row_offsets,
        segment_lengths, output_ptr,
        A_gather_locations, B_gather_locations,
        I, J, V);

    //slices.push_back(Container());
    //slices.back().swap(C_slice);

    begin_row = end_row;

    // at this point, we have the slice of A[slice,:]*B
    // need to multiply this slice with C to get R[slice,:]
    // no need to split this call - temp_slice should be small(ish)
    // compute row offsets for C

    cusp::array1d<IndexType,MemorySpace> C_row_offsets(C.num_rows + 1);
    cusp::detail::indices_to_offsets(C.row_indices, C_row_offsets);

    // compute row lengths for C
    cusp::array1d<IndexType,MemorySpace> C_row_lengths(C.num_rows);
    thrust::transform(C_row_offsets.begin() + 1, C_row_offsets.end(), C_row_offsets.begin(), C_row_lengths.begin(), \
        thrust::minus<IndexType>());

    // for each element temp_slice(i,j) compute the number of nonzero elements in C(j,:)
    cusp::array1d<IndexType,MemorySpace> C_segment_lengths(temp_slice.num_entries);
    thrust::gather(temp_slice.column_indices.begin(), temp_slice.column_indices.end(),
        C_row_lengths.begin(),
        C_segment_lengths.begin());

    // output pointer
    cusp::array1d<IndexType,MemorySpace> C_output_ptr(temp_slice.num_entries + 1);
    thrust::exclusive_scan(C_segment_lengths.begin(), C_segment_lengths.end(),
        C_output_ptr.begin(),
        IndexType(0));
    C_output_ptr[temp_slice.num_entries] = C_output_ptr[temp_slice.num_entries - 1] + C_segment_lengths[temp_slice.num_entries - 1]; // XXX is this necessary?

    // number of non-zeros in result?
    size_t C_coo_num_nonzeros = C_output_ptr[temp_slice.num_entries];

    C_begin_row = 0;
    C_end_row = temp_slice.num_rows;
    C_begin_segment = 0;
    C_end_segment = temp_slice.num_entries;

    size_t C_workspace_size = C_output_ptr[C_end_segment] - C_output_ptr[C_begin_segment];

    coo_spmm_helper(C_workspace_size,
        C_begin_row, C_end_row,
        C_begin_segment, C_end_segment,
        temp_slice, C, R_slice,
        C_row_offsets,
        C_segment_lengths, C_output_ptr,
        temp_gather_locations, C_gather_locations,
        C_I, C_J, C_V);

    // instead of keeping slices, add directly to output matrix?
    slices.push_back(Container());
    slices.back().swap(R_slice);

    call_no++;
  } // end while

  // deallocate workspace
  A_gather_locations.clear(); A_gather_locations.shrink_to_fit();  
  B_gather_locations.clear(); B_gather_locations.shrink_to_fit();

  I.clear();   I.shrink_to_fit();
  J.clear();   J.shrink_to_fit();
  V.clear();   V.shrink_to_fit();

  temp_gather_locations.clear(); temp_gather_locations.shrink_to_fit();
  C_gather_locations.clear(); C_gather_locations.shrink_to_fit();
  C_I.clear(); C_I.shrink_to_fit();
  C_J.clear(); C_J.shrink_to_fit();
  C_V.clear(); C_V.shrink_to_fit();

  // compute total output size
  size_t R_num_entries = 0;
  for(typename ContainerList::iterator iter = slices.begin(); iter != slices.end(); ++iter) {
    R_num_entries += iter->num_entries;
  }
  // resize output
  D.resize(A.num_rows, C.num_cols, R_num_entries);

  // copy slices into output
  // this seems unnecessary - keeps 2 copies of the output in memory...
  // delete slice as soon as copied?
  size_t base = 0;
  for(typename ContainerList::iterator iter = slices.begin(); iter != slices.end(); ++iter)
  {
    thrust::copy(iter->row_indices.begin(),    iter->row_indices.end(),    D.row_indices.begin()    + base);
    thrust::copy(iter->column_indices.begin(), iter->column_indices.end(), D.column_indices.begin() + base);
    thrust::copy(iter->values.begin(),         iter->values.end(),         D.values.begin()         + base);
    base += iter->num_entries;
    // manually delete container entry?
    // (*iter)->Container::~Container();
  } // end for
}

} // end namespace device
} // end namespace detail
} // end namespace cusp
